import numpy as np
import xmltodict as xtd
import numpy as np

def read_ecg(fname):

    # names and order of leads to import
    lead_order = ["I", "II", "III", "AVR", "AVL", "AVF", "V1", "V2", "V3", "V4", "V5", "V6"]

    ecg, md = import_ecg(fname, lead_order)

    return ecg, md

def import_ecg(fname, lead_order):
    f = open(fname, "rt")
    raw_input = f.read()
    f.close()

    # extract ECG and metadata
    ecg, md = parse_xml(raw_input, lead_order)
    md['filename'] = fname

    return ecg, md


def parse_xml(input_data, lead_order):
    """ Takes input as raw xml data read from file (Cardiosoft specification), returns list of 12 lead waveforms
    and metadata. """
    md = {}  # metadata dictionary
    data = xtd.parse(input_data)['CardiologyXML']

    full_lead_nodes = [['StripData', 'WaveformData'], ['Strip', 'StripData', 'WaveformData']]
    median_lead_nodes = [['RestingECGMeasurements', 'MedianSamples', 'WaveformData'], []]
    full_leads = get_lead_data(data, full_lead_nodes, lead_order)
    # median_leads = get_lead_data(data, median_lead_nodes, lead_order)

    md = get_metadata(data, md)
    md['lead order'] = lead_order

    return np.array(full_leads, dtype=np.int16), md


def get_lead_data(data, nodes, lead_order):
    leads = [[] for i in range(12)]
    raw_lead_data = get_xml_node(data, nodes)

    if len(raw_lead_data) != 12:
        print('Warning: only {} leads found'.format(len(raw_lead_data)))

    else:
        for i in range(len(raw_lead_data)):
            lead_n = lead_order.index(raw_lead_data[i]['@lead'].upper())
            lead_data_string = raw_lead_data[i]['#text']
            lead_vals = [int(x) for x in lead_data_string.split(",")]
            leads[lead_n] = lead_vals

    return leads


def get_xml_node(data, node_list):
    output = data
    try:
        for x in node_list[0]:
            output = output[x]
    except KeyError:
        output = data[:]
        try:
            for x in node_list[1]:
                output = output[x]
        except:
            raise ValueError('No lead data found!')
    return output


def safe_extract(data, keys, default='N/A', transform=None):
    """
    Safely extract a value from a nested dictionary using a list of keys.
    If any key is missing, or a transformation fails, return a default value.
    
    Parameters:
    - data: The dictionary from which to extract the value.
    - keys: A list of keys to navigate through the dictionary.
    - default: The default value to return if extraction fails.
    - transform: An optional function to transform the extracted value.
    
    Returns:
    - The extracted and possibly transformed value, or the default value.
    """
    try:
        # Navigate through the nested dictionary using the keys
        for key in keys:
            data = data[key]
        # Apply transformation function if provided
        if transform:
            return transform(data)
        return data
    except (KeyError, TypeError, ValueError):
        return default

def get_metadata(data, md):
    # Use safe_extract to attempt extraction and set to default if it fails
    md['sample rate'] = safe_extract(data, ['StripData', 'SampleRate', '#text'], default='N/A', transform=float)
    md['t scale'] = 1. / float(md['sample rate']) if md['sample rate'] != 'N/A' else 'N/A'
    md['v scale'] = safe_extract(data, ['StripData', 'Resolution', '#text'], default='N/A', transform=float)
    md['filter 50Hz'] = safe_extract(data, ['FilterSetting', 'Filter50Hz'], default='N/A')
    md['filter 60Hz'] = safe_extract(data, ['FilterSetting', 'Filter60Hz'], default='N/A')
    md['low pass'] = safe_extract(data, ['FilterSetting', 'LowPass', '#text'], default='N/A', transform=float)
    md['high pass'] = safe_extract(data, ['FilterSetting', 'HighPass', '#text'], default='N/A', transform=float)

    meas = data.get('RestingECGMeasurements', {})
    md['Heart rate'] = safe_extract(meas, ['VentricularRate', '#text'], default='N/A') + ' ' + safe_extract(meas, ['VentricularRate', '@units'], default='')
    md['P duration'] = safe_extract(meas, ['PDuration', '#text'], default='N/A') + ' ' + safe_extract(meas, ['PDuration', '@units'], default='')
    md['PR interval'] = safe_extract(meas, ['PQInterval', '#text'], default='N/A') + ' ' + safe_extract(meas, ['PQInterval', '@units'], default='')
    md['QRS duration'] = safe_extract(meas, ['QRSDuration', '#text'], default='N/A') + ' ' + safe_extract(meas, ['QRSDuration', '@units'], default='')
    md['QT interval'] = safe_extract(meas, ['QTInterval', '#text'], default='N/A') + ' ' + safe_extract(meas, ['QTInterval', '@units'], default='')
    md['QTc interval'] = safe_extract(meas, ['QTCInterval', '#text'], default='N/A') + ' ' + safe_extract(meas, ['QTCInterval', '@units'], default='')
    md['P axis'] = safe_extract(meas, ['PAxis', '#text'], default='N/A') + ' ' + safe_extract(meas, ['PAxis', '@units'], default='')
    md['R axis'] = safe_extract(meas, ['RAxis', '#text'], default='N/A') + ' ' + safe_extract(meas, ['RAxis', '@units'], default='')
    md['T axis'] = safe_extract(meas, ['TAxis', '#text'], default='N/A') + ' ' + safe_extract(meas, ['TAxis', '@units'], default='')
    return md
